'''Caesar Cipher'''

'''ex1'''

'''
start
    alphabet = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
    shifted_alphabet = alphabet[shift:] + alphabet[:shift]
    table = str.maketrans(alphabet, shifted_alphabet)
    return text.translate(table)
end
'''

# def caesar_cipher(text, shift):
#     alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
#     shifted_alphabet = alphabet[shift:] + alphabet[:shift]
#     table = str.maketrans(alphabet, shifted_alphabet)
#     return text.translate(table)

# message = "HELLO WORLD"
# shift = 20
# encrypted = caesar_cipher(message, shift)
# print(f"Original: {message}")
# print(f"Encrypted: {encrypted}")

# def decrypt_caesar_cipher(encrypted, shift):
#     return caesar_cipher(encrypted, -shift)

# decrypted = decrypt_caesar_cipher(encrypted, shift)
# print(f"Decrypted: {decrypted}")


# print("-" * 100)

'''ex2'''

'''Vigenère Cipher'''

'''
message = "HELLO WORLD"
keyword = "KEY"
encrypted = vigenere_cipher(message, keyword)
decrypted = vigenere_cipher(encrypted, keyword, decrypt=True)
'''
'''
start
alphabet = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
key_length = len(key)
message_length = len(message)
if key_length < message_length:
    key = (key * (message_length // key_length + 1))[:message_length]


end
'''

def vigenere_cipher(message, key, decrypt=False):
    alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
    key_index = 0
    result = ""

    key_length = len(key) # Ensure the key is as long as the message
    message_length = len(message)
    if key_length < message_length:
        key = (key * (message_length // key_length + 1))[:message_length]
    
    
    for char in message:
        if char in alphabet:
            char_pos = alphabet.find(char)
            key_shift = alphabet.find(key[key_index])
            
            if decrypt:
                result += alphabet[(char_pos - key_shift) % len(alphabet)]
            else:
                result += alphabet[(char_pos + key_shift) % len(alphabet)]
            
            key_index = (key_index + 1) % key_length
        else:
            # Non-alphabet characters are added unchanged
            result += char
            
    return result

# # Test the function
# message = "HELLO WORLD"
# keyword = "KEY"
# encrypted = vigenere_cipher(message, keyword)
# decrypted = vigenere_cipher(encrypted, keyword, decrypt=True)

# print(f"Original: {message}")
# print(f"Encrypted: {encrypted}")
# print(f"Decrypted: {decrypted}")
    

